"""
使用 YOLO 模型，预测摄像头输入，测试推理结果

date: 2025/9/27
author: SiHeng Tang
file: predict_test.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import argparse
import base64
import logging

import cv2
import numpy as np

from ultralytics import YOLO

logger = logging.getLogger(__name__)


def _encode_frame(f: np.ndarray, q: int) -> str:
    """
    WebP 有损编码，返回字节
    """
    ok, buffer = cv2.imencode('.webp', f, [cv2.IMWRITE_WEBP_QUALITY, q])

    if not ok:
        raise RuntimeError("image encode failed")

    return base64.b64encode(buffer.tobytes()).decode('ascii')


def _decode_frame(encoded_bytes: str) -> np.ndarray:
    """
    解码 WebP 图像字节为 numpy 数组（BGR 格式）
    """
    raw_array = np.frombuffer(base64.b64decode(encoded_bytes), np.uint8)
    cv_image = cv2.imdecode(raw_array, cv2.IMREAD_COLOR)

    if cv_image is None:
        raise RuntimeError("image decode failed")

    return cv_image


def crop_rotated_obb(img: np.ndarray, xywhr: np.ndarray) -> np.ndarray:
    """
    从原图中抠出旋转框并旋转至水平（保持比例）
    :param img: 原图 BGR
    :param xywhr: 旋转框参数 (x, y, w, h, r)，其中r为旋转角度(弧度)
    :return: 旋转并裁剪后的小图
    """
    # 解析xywhr参数，转换弧度为角度
    cx, cy, w, h, r = xywhr
    angle = np.rad2deg(r)

    if w < h:
        w, h = h, w
        angle += 90

    # 旋转整张图
    m = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
    rows, cols = img.shape[:2]
    rotated_img = cv2.warpAffine(img, m, (cols, rows), borderValue=(0, 0, 0))

    # 裁剪
    x1 = max(0, int(round(cx - w / 2)))
    y1 = max(0, int(round(cy - h / 2)))
    x2 = min(cols, int(round(cx + w / 2)))
    y2 = min(rows, int(round(cy + h / 2)))

    cropped = rotated_img[y1:y2, x1:x2]
    return cropped


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='YOLO OBB Detection Program')
    parser.add_argument('cam_path', help='Camera device path')
    parser.add_argument('model_path', help='Model file path')
    parser.add_argument('--quality', type=int, default=100, help='Image encoding quality (default: 100)')
    parser.add_argument('--show-image', action='store_true', default=False,
                        help='Whether to show image in OpenCV (default: False)')
    parser.add_argument('--frame-shape', type=int, nargs=2, default=(1280, 800),
                        help='Frame size (width height) (default: 1280 800)')
    parser.add_argument('--cap-fps', type=int, default=120, help='Capture frame rate (default: 120)')

    args = parser.parse_args()

    cam_path = args.cam_path
    model_path = args.model_path
    quality = args.quality
    show_image_in_cv = args.show_image
    frame_shape = tuple(args.frame_shape)
    cap_fps = args.cap_fps

    cap = cv2.VideoCapture(cam_path)
    model = YOLO(model_path)

    if not cap.isOpened():
        logger.error(f"cannot open camera {cam_path}")
        exit(1)

    cap.set(cv2.CAP_PROP_FRAME_WIDTH, frame_shape[0])
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, frame_shape[1])
    cap.set(cv2.CAP_PROP_FPS, cap_fps)

    try:
        while True:
            ret, frame = cap.read()

            if not ret:
                logger.error("camera read failed, retry")
                continue

            str_full_scene_img = _encode_frame(frame, quality)

            # Run inference on the source
            results = model(frame)
            for result in results:
                names = [result.names[cls.item()] for cls in result.obb.cls.int()]

                for i, (pts, name, conf) in enumerate(zip(result.obb.xyxyxyxy, names, result.obb.conf)):
                    patch = crop_rotated_obb(frame, result.obb.xywhr[i].cpu().numpy())

                    str_target_img = _encode_frame(patch, quality)
                    print(f"Target slice image {result.obb.xywhr=} {result.obb.xyxyxyxy=} "
                          f"{name=} {i=} {str_target_img=}")

                    if show_image_in_cv:
                        cv2.imshow(f"{name=} {i=}", _decode_frame(str_target_img))
                        cv2.waitKey(0)

            print(f"Full scene image {str_full_scene_img}")

    except KeyboardInterrupt:
        logger.error("keyboard interrupt")
        cap.release()
        exit(0)
